package com.edu.framework.websocket;

import com.alibaba.fastjson2.JSONException;
import com.alibaba.fastjson2.JSONObject;
import com.edu.common.utils.StringUtils;
import com.edu.framework.websocket.message.base.Message;
import java.util.List;
import java.util.function.Consumer;
import java.util.function.Function;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * websocket 客户端用户集
 *
 * @author xxxx
 */
@Slf4j
public class WebSocketHolder {

    private static final Map<String, WebSocketSession> CODE_SESSION_MAP = new ConcurrentHashMap<>();

    private static final Map<WebSocketSession, String> SESSION_CODE_MAP = new ConcurrentHashMap<>();


    public static WebSocketSession getSessionByCode(String code){
        return CODE_SESSION_MAP.get(code);
    }

    public static Map<String, WebSocketSession> getSessions(){
        return CODE_SESSION_MAP;
    }

    public static String getCodeBySession(WebSocketSession session){
        return SESSION_CODE_MAP.get(session);
    }
    public static void addSession(WebSocketSession session, String terminalCode) {
        if(CODE_SESSION_MAP.containsKey(terminalCode)){
            WebSocketSession webSocketSession = CODE_SESSION_MAP.get(terminalCode);
            SESSION_CODE_MAP.remove(webSocketSession);
        }
        CODE_SESSION_MAP.put(terminalCode, session);
        SESSION_CODE_MAP.put(session, terminalCode);
    }

    public static void removeSession(WebSocketSession session) {
        String terminalCode = SESSION_CODE_MAP.remove(session);
        if (StringUtils.isNotEmpty(terminalCode)) {
            CODE_SESSION_MAP.remove(terminalCode);
        }
    }

    public static <T extends Message> void send(WebSocketSession session, T message) throws JSONException, IOException {
        String messageText = buildTextMessage(message);
        sendTextMessage(session, messageText);
    }

    public static <T extends Message> boolean send(String terminalCode, T message) throws JSONException, IOException {
        WebSocketSession session = CODE_SESSION_MAP.get(terminalCode);
        if (session == null) {
            return false;
        }
        send(session, message);
        return true;
    }


    private static <T extends Message> String buildTextMessage(T message) throws JSONException {
        return JSONObject.toJSONString(message);
    }

    private static void sendTextMessage(WebSocketSession session, String messageText) throws IOException {
        if (session == null || messageText == null) {
            return;
        }
        session.sendMessage(new TextMessage(messageText));
    }

    public static void sendBinaryMessage(WebSocketSession session, BinaryMessage message) throws IOException {
        if (session == null) {
            return;
        }
        session.sendMessage(message);
    }

    public static <T extends Message> void broadcast(T message) throws IOException, JSONException {
        String messageText = buildTextMessage(message);
        if (messageText == null){
            return;
        }
        for (WebSocketSession session : SESSION_CODE_MAP.keySet()) {
            sendTextMessage(session, messageText);
        }
    }

    public static <T extends Message> void broadcast(Function<String, Message> function) throws IOException {
        for (WebSocketSession session : SESSION_CODE_MAP.keySet()) {
            String code = (String) session.getAttributes().get("code");
            Message apply = function.apply(code);
            if (apply == null){
                continue;
            }
            String messageText = buildTextMessage(apply);

            sendTextMessage(session, messageText);
        }
    }

    public static <T extends Message> void broadcast(Function<String, Message> function, Consumer<Message> invoke) throws IOException {
        for (WebSocketSession session : SESSION_CODE_MAP.keySet()) {
            String code = (String) session.getAttributes().get("code");
            Message apply = function.apply(code);
            if (apply == null){
                continue;
            }
            String messageText = buildTextMessage(apply);
            sendTextMessage(session, messageText);

            // 发布完成后 处理逻辑
            if (invoke != null){
                invoke.accept(apply);
            }
        }
    }

    public static <T extends Message> void broadcastForCode(List<String> terminalCodes, Function<String, Message> function, Consumer<Message> invoke) throws IOException {
        for (String code : terminalCodes){
            WebSocketSession session = CODE_SESSION_MAP.get(code);
            if (session != null){
                Message apply = function.apply(code);
                sendTextMessage(session, buildTextMessage(apply));
                // 发布完成后 处理逻辑
                if (invoke != null){
                    invoke.accept(apply);
                }
            }
        }
    }







}
